import os
import numpy as np
import torch
from PIL import Image


class GraspDataset(object):
    def __init__(self, root, transforms, transforms1):
        self.root = root
        self.transforms = transforms
        self.transforms1 = transforms1
        # load all image files, sorting them to
        # ensure that they are aligned
        self.imgs = list(sorted(os.listdir(os.path.join(root, "Images"))))
        self.masks = list(sorted(os.listdir(os.path.join(root, "Masks"))))

    def __getitem__(self, idx):
        # load images and masks
        img_path = os.path.join(self.root, "Images", self.imgs[idx])
        mask_path = os.path.join(self.root, "Masks", self.masks[idx])

        img = Image.open(img_path).convert("RGB")

        # img = np.load(img_path)
        
        
        # img = Image.fromarray(img).convert("RGB")
        

        # note that we haven't converted the mask to RGB,
        # because each color corresponds to a different instance
        # with 0 being background

        mask = Image.open(mask_path)

        # mask = np.load(mask_path)

        # mask = Image.fromarray(mask)

        # convert the PIL Image into a numpy array
        mask = np.array(mask)
        # 取颜色不同的实例对象
        obj_ids = np.unique(mask)
        # first id is the background, so remove it
        obj_ids = obj_ids[1:]

        # split the color-encoded mask into a set
        # of binary masks
        masks = mask == obj_ids[:, None, None]

        # 获取每个mask对应的bbox的坐标
        # num_objs = len(obj_ids)
        # boxes = []
        # for i in range(num_objs):
        #     pos = np.where(masks[i])
        #     # pos[1]、pos[0]包含mask上的全部x坐标、y坐标
        #     # 通过获取mask上的最大和最小的x、y坐标，来获取bbox
        #     xmin = np.min(pos[1])
        #     xmax = np.max(pos[1])
        #     ymin = np.min(pos[0])
        #     ymax = np.max(pos[0])
        #     boxes.append([xmin, ymin, xmax, ymax])

        # # convert everything into a torch.Tensor
        # boxes = torch.as_tensor(boxes, dtype=torch.float32)
        # # there is only one class
        # labels = torch.ones((num_objs,), dtype=torch.int64)
        # masks = torch.as_tensor(masks, dtype=torch.uint8)

        # image_id = torch.tensor([idx])
        # # 计算bbox面积
        # area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        # # suppose all instances are not crowd
        # iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        target = {}
        # target["boxes"] = boxes
        # target["labels"] = labels
        # target["masks"] = masks
        # target["image_id"] = image_id
        # target["area"] = area
        # target["iscrowd"] = iscrowd

        if self.transforms is not None:
            img = self.transforms(img)
            # target = self.transforms(target)   target是一个字典，包含标注文件中的各种信息，不能用pil处理
            
            mask = self.transforms1(mask)

            # img, target = self.transforms(img, target)

        # return img, mask
        return img, target, mask

    def __len__(self):
        return len(self.imgs)